#!/usr/bin/env python
#-*- coding: UTF-8 -*-
# schif.py: node manager interface functions for scheduler
#
# Copyright (C) 2008-2010 Hily Jiang <hilyjiang@gmail.com>
#
# This file is part of dtask.
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

"""
Node Manager interface functions for scheduler.

Scheduler <-> Node Manager Protocol:

* New Connection
  (S -> NM) Send data format:    PASSPORT schName schKey\n

* Dispatch Task
  (S -> NM) Send data format:    TASK taskId taskType taskParams\n

* Return Result
  (NM -> S) Send data format:    RESULT taskId result\n

* Update Task Status
  (NM -> S) Send data format:    STATUS taskId status\n
"""

import sys
import threading

from time import time

from utils import CommandBuffer


CFG_SCH_BUFSIZ = 1024


class SchedulerIfBuffer(CommandBuffer):
    """Scheduler buffer to receive command string and scheduler info."""

    __schedulerCount = 0

    def __init__(self):
        CommandBuffer.__init__(self)
        SchedulerIfBuffer.__schedulerCount += 1
        self.schId = SchedulerIfBuffer.__schedulerCount
        self.schName = ''
        self.schKey = ''
        self.hasInfo = False

    def setInfo(self, schName='', schKey=''):
        """Set scheduler info."""
        if self.hasInfo:
            return False
        self.schName = schName
        self.schKey = schKey
        self.hasInfo = True

    def getInfo(self):
        """Get scheduler info."""
        return (self.schName, self.schKey)

    def getSchedulerId(self):
        return self.schId


def schCommandHandler(cmdInfo, nodeManager, dispatcher, siBuf,
                      schClientSock, schClientName,
                      nodSockDict, schSockDict):
    """Handle the scheduler's commands."""
    cmd, params = cmdInfo

    print schClientName, '[SCH]', cmd, 
    if params:
        print params[:40]
    else:
        print

    if cmd == 'TASK':
        try:
            taskId, taskType, taskParams = params.split(None, 2)
        except Exception, e:
            print schClientName, '[SCH]', e
        else:
            nodeId = nodeManager.acquireNode()
            schId = siBuf.getSchedulerId()
            dispatcher.dispatch(schId, taskId, nodeId)
            newTaskId = '%s|%s' % (schId, taskId) # 'schid|orginalTaskId'
            if nodSockDict.has_key(nodeId):
                nodeSock = nodSockDict[nodeId]
                nodeSock.sendall('TASK %s %s %s\n' % 
                                 (newTaskId, taskType, taskParams))
                nodName, nodIp, nodPort = nodeManager.getNodeInfo(nodeId)
                schClientSock.sendall('STATUS %s %s %s %s %s 0\n' %
                         (taskId, nodName, nodIp, nodPort, int(time())))
    elif cmd == 'PASSPORT':
        params = params.split(None, 1)
        if len(params) == 1:
            schType = params[0]
            schName = schKey = ''
        else:
            schName, schKey = params
        siBuf.setInfo(schName, schKey)
        schId = siBuf.getSchedulerId()
        schSockDict[schId] = schClientSock
        # todo: authorization


def schRecvData(nodeManager, dispatcher, siBuf,
                schClientSock, schClientName, data,
                nodSockDict, schSockDict):
    """Process data send by scheduler."""
    siBuf.addData(data)

    while True:
        cmdInfo = siBuf.getCommand()
        if cmdInfo:
            schCommandHandler(cmdInfo, nodeManager, dispatcher, siBuf,
                              schClientSock, schClientName,
                              nodSockDict, schSockDict)
        else:
            break


def schRecvThreadFunc(nodeManager, dispatcher, schClientSock, schClientSocks,
                      nodSockDict, schSockDict):
    """Thread function for each scheduler connection to process data."""
    schClientName = schClientSock.getpeername()
    siBuf = SchedulerIfBuffer()

    while True:
#        try:
            data = schClientSock.recv(CFG_SCH_BUFSIZ)
            if data == '': break
            schRecvData(nodeManager, dispatcher, siBuf,
                        schClientSock, schClientName, data,
                        nodSockDict, schSockDict)
#        except Exception, e:
#            print schClientName, '[SCH]', sys._getframe().f_lineno, e
    
    print schClientName, '[SCH] quit'
    schId = siBuf.getSchedulerId()
    if schSockDict.has_key(schId):
        del schSockDict[schId]
    schClientSock.close()
    schClientSocks.remove(schClientSock)


def startSchedulerIf(nodeManager, dispatcher, schServerSock, schClientSocks,
                     nodSockDict, schSockDict):
    """Start interface for schedulers."""
    schRecvThreads = []

    while True:
        schClientSock, schClientAddr = schServerSock.accept()
        schClientSocks.append(schClientSock)
        print schClientAddr, '[SCH] connected'
        t = threading.Thread(target=schRecvThreadFunc,
                   args=(nodeManager, dispatcher, schClientSock,
                         schClientSocks, nodSockDict, schSockDict))
        t.setDaemon(True)
        schRecvThreads.append(t)
        t.start()

    print 'Waiting for schedulers to release...',
    for t in schRecvThreads:
        t.join()
    print 'ok'
